package com.zw.cloud.tools.controller.ocr;

import com.alibaba.fastjson2.JSON;
import com.baidu.aip.ocr.AipOcr;
import com.zw.cloud.tools.modle.vo.OcrExcelVO;
import com.zw.cloud.tools.modle.vo.OcrVO;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.poi.xssf.streaming.SXSSFRow;
import org.apache.poi.xssf.streaming.SXSSFSheet;
import org.apache.poi.xssf.streaming.SXSSFWorkbook;
import org.json.JSONObject;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.ResourceUtils;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;

import javax.servlet.http.HttpServletResponse;
import java.io.File;
import java.io.FileOutputStream;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.List;
import java.util.Objects;

/**
 * 文字识别
 * https://ai.baidu.com/ai-doc/OCR/Nkibizxlf
 */
@RequestMapping("/ocr")
@RestController
@Slf4j
public class OcrController {

    @Autowired
    private AipOcr client;

    /**
     * 高精度识别
     */
    @PostMapping("/highPrecisionImageRecognition")
    //http://localhost:9040/ocr/highPrecisionImageRecognition
    public OcrVO highPrecisionImageRecognition(MultipartFile file) throws Exception{
        // 传入可选参数调用接口
        HashMap<String, String> options = buildOptionsHashMap();
        // 参数为本地图片路径
        /*String image = "test.jpg";
        JSONObject res = client.basicAccurateGeneral(image, options);
        System.out.println(res.toString(2));*/

        // 通用文字识别（含生僻字版）, 图片参数为远程url图片
       /* JSONObject res = client.enhancedGeneralUrl(url, options);
        System.out.println(res.toString(2));*/

        // 参数为本地图片二进制数组
        byte[] bytes = file.getBytes();
        JSONObject res = client.basicAccurateGeneral(bytes, options);
        //System.out.println(res.toString(2));
        return JSON.parseObject(JSON.toJSONString(res), OcrVO.class);
    }

    /**
     * 网络图片文字识别
     */
    @PostMapping("/netImageRecognition")
    //http://localhost:9040/ocr/netImageRecognition
    public OcrVO netImageRecognition(MultipartFile file) throws Exception{
        HashMap<String, String> options = buildOptionsHashMap();
        /* // 参数为本地图片路径
        String image = "test.jpg";
        JSONObject res = client.webImage(image, options);
        System.out.println(res.toString(2));*/

        // 网络图片文字识别, 图片参数为远程url图片
        /*JSONObject res = client.webImageUrl(url, options);
        System.out.println(res.toString(2));*/

        // 参数为本地图片二进制数组
        byte[] bytes = file.getBytes();
        JSONObject res = client.basicAccurateGeneral(bytes, options);
        return JSON.parseObject(JSON.toJSONString(res), OcrVO.class);
    }

    /**
     * 身份证识别
     */
    @PostMapping("/icCardRecognition")
    //http://localhost:9040/ocr/icCardRecognition
    public OcrVO icCardRecognition(MultipartFile file) throws Exception{
        HashMap<String, String> options = buildOptionsHashMap();
        String idCardSide = "back";

        // 参数为本地图片路径
       /* String image = "test.jpg";
        JSONObject res = client.idcard(image, idCardSide, options);
        System.out.println(res.toString(2));*/

        byte[] bytes = file.getBytes();
        JSONObject res = client.idcard(bytes, idCardSide, options);
        return JSON.parseObject(JSON.toJSONString(res), OcrVO.class);
    }

    /**
     * 手写文字识别
     */
    @PostMapping("/handwritingRecognition")
    //http://localhost:9040/ocr/handwritingRecognition
    public OcrVO handwritingRecognition(MultipartFile file) throws Exception{
        // 传入可选参数调用接口
        HashMap<String, String> options = buildOptionsHashMap();
        options.put("recognize_granularity","big");
        // 参数图片url
        /*String url = "http://localhost/test.jpg"
        res = client.handwritingUrl(url, options);
        System.out.println(res.toString(2));*/

        JSONObject res = client.handwriting(file.getBytes(), options);
        return JSON.parseObject(JSON.toJSONString(res), OcrVO.class);
    }

    /**
     * 二维码识别
     */
    @PostMapping("/codeRecognition")
    //http://localhost:9040/ocr/codeRecognition
    public String codeRecognition(MultipartFile file) throws Exception{
        // 传入可选参数调用接口
        HashMap<String, String> options = new HashMap<String, String>();

        // 参数图片url
       /* String url = "http://localhost/test.jpg"
        res = client.qrcodeUrl(url, options);
        System.out.println(res.toString(2));
        byte[] bytes = file.getBytes();*/
        JSONObject res = client.qrcode(file.getBytes(), options);
        return res.toString();
    }

    /**
     * 表格文字识别转为excel
     */
    @PostMapping("/formRecognition")
    //http://localhost:9040/ocr/formRecognition
    public void formRecognition(MultipartFile file, Integer source, String fileName, HttpServletResponse response) throws Exception {
        log.info("[OcrController][formRecognition]source is {},fileName is {}", source, fileName);
        // 传入可选参数调用接口
        HashMap<String, String> options = new HashMap<>();

        JSONObject res = client.form(file.getBytes(), options);
        //System.out.println(res.toString(2));
        List<OcrExcelVO> ocrExcelVOS = JSON.parseArray((res.getJSONArray("forms_result").toString()), OcrExcelVO.class);
        System.out.println(JSON.toJSONString(ocrExcelVOS));
        exportExcel(ocrExcelVOS, source, fileName, response);
    }

    private void exportExcel(List<OcrExcelVO> ocrExcelVOList, Integer source, String fileName, HttpServletResponse response) throws Exception {
        //创建HSSFWorkbook对象(excel的文档对象)
        SXSSFWorkbook workbook = new SXSSFWorkbook(200);
        //建立新的sheet对象（excel的表单）
        SXSSFSheet sheet = workbook.createSheet("用户数据");
        //在sheet里创建第一行，参数为行索引(excel的行)
        int rowNum = 0;
        for (OcrExcelVO ocrExcelVO : ocrExcelVOList) {
            List<OcrExcelVO.KeyWord> header = ocrExcelVO.getHeader();
            rowNum = fillCellValue(header, sheet, rowNum);

            List<OcrExcelVO.KeyWord> body = ocrExcelVO.getBody();
            rowNum = fillCellValue(body, sheet, rowNum);

            List<OcrExcelVO.KeyWord> footer = ocrExcelVO.getFooter();

            rowNum = fillCellValue(footer, sheet, rowNum);

            rowNum++;
        }

        if (Objects.equals(1, source)) {
            if (StringUtils.isBlank(fileName)) {
                throw new RuntimeException("fileName is null");
            }
            String absolutePath = ResourceUtils.getFile("classpath:").getAbsolutePath();
            File file = new File(absolutePath + "/excels/" + fileName);
            if (!file.getParentFile().exists()) {
                file.getParentFile().mkdirs();
            }
            FileOutputStream stream = new FileOutputStream(file);
            workbook.write(stream);
            stream.close();
        } else {
            response.setCharacterEncoding("UTF-8");
            response.setHeader("content-Type", "application/vnd.ms-excel");
            response.setHeader("Content-Disposition",
                    "attachment;filename=" + URLEncoder.encode(fileName, String.valueOf(StandardCharsets.UTF_8)));
            workbook.write(response.getOutputStream());
        }

    }

    private int fillCellValue(List<OcrExcelVO.KeyWord> keyWordList, SXSSFSheet sheet, int rowNum) {
        if (CollectionUtils.isEmpty(keyWordList)) {
            return rowNum;
        }
        Integer currRowIndex = 0;
        SXSSFRow row = sheet.createRow(rowNum++);
        for (OcrExcelVO.KeyWord keyWord : keyWordList) {
            Integer rowIndex = keyWord.getRow();

            if (!currRowIndex.equals(rowIndex)) {
                row = sheet.createRow(rowNum);
                rowNum++;
                currRowIndex++;
            }

            row.createCell(keyWord.getColumn()).setCellValue(keyWord.getWords());

        }
        return rowNum;
    }


    private HashMap<String, String> buildOptionsHashMap() {
        // 传入可选参数调用接口
        HashMap<String, String> options = new HashMap<String, String>(8);
        options.put("language_type", "CHN_ENG");

        //是否检测图像朝向，
        options.put("detect_direction", "false");
        //是否返回识别结果中每一行的置信度
        options.put("probability", "false");
        return options;
    }

}
